import concurrent.futures
import queue
import subprocess
import threading
import time

import psutil

from config.LoadConfig import get_config
from utils.Singleton import Singleton

system_conf = get_config("system")


class DeviceChecker:
    @classmethod
    def get_runtime_env_device(cls):
        """获取运行环境"""
        try:
            runtime_env = system_conf.get('runtime_env', None)
            return 'gpu' if runtime_env == 'gpu' else 'cpu'
        except:
            pass
        # 判断是否有nvidia驱动
        if DeviceChecker.check_nvidia_driver():
            return 'gpu'
        else:
            return 'cpu'

    @classmethod
    def check_nvidia_driver(cls):
        try:
            output = subprocess.check_output("nvidia-smi", shell=True)
            print("NVIDIA GPU 设备已检测到")
            return True
        except FileNotFoundError:
            print("未安装 NVIDIA GPU 或驱动")
            return False
        except subprocess.CalledProcessError:
            print("nvidia-smi 命令无法执行，可能未正确安装驱动")
            return False
        except Exception as e:
            print(e)
            return False

    @classmethod
    def check_cpu_availability(cls):
        """
        检查CPU余量 以及内存剩余 MB
        """
        cpu_availability = psutil.cpu_percent()  # 当前CPU使用率，比如 60%
        memory_available = psutil.virtual_memory().available / 1024 / 1024  # 剩余内存（MB）
        return 100 - cpu_availability, memory_available # 剩余40% ，可用内存

    @classmethod
    def check_gpu_availability(cls):
        """
        检查GPU显存余量
        :return 剩余GPU使用率 显存余量 MB
        """
        output = subprocess.check_output(
            "nvidia-smi --query-gpu=utilization.gpu,memory.total,memory.used,memory.free --format=csv,noheader,nounits",
            shell=True
        ).decode("utf-8").strip()
        total_free_mem = 0
        avg_gpu_used = 0
        device_count = 0
        for line in output.splitlines():
            output = line
            device_count += 1
            gpu_used, total_mem, used_mem, free_mem = map(int, output.split(", "))
            total_free_mem += int(free_mem)
            avg_gpu_used += gpu_used
        return 100 - (avg_gpu_used / device_count), total_free_mem

class JobWatchDog:
    """
    任务监控狗
    """
    def __init__(self, device:str, cpu_availability_limit:float, memory_limit:float, gpu_availability_limit:float, gpu_memory_limit:float, min_available_threads_num:int):
        # 判断有没有gpu
        self.device: str = device
        # CPU使用率阈值
        self.cpu_availability_limit = cpu_availability_limit
        # 内存阈值
        self.memory_limit = memory_limit
        # GPU使用率阈值
        self.gpu_availability_limit = gpu_availability_limit
        # GPU内存使用阈值
        self.gpu_memory_limit = gpu_memory_limit
        self.threads_num_locker = threading.Lock()
        # 最低可用线程数
        self.min_available_threads_num = min_available_threads_num

    def get_available_threads_num(self, need_threads_num, used_threads_num, total_threads_num):
        """
        通过处理器资源 计算可新增线程数
        当处理器资源不足并且没有已使用线程时 需要保证最低线程数
        """
        cpu_availability, memory_available = DeviceChecker.check_cpu_availability()
        # 计算空余CPU资源
        """
        假设
        system预设值:
            cpu_availability_limit: 20    # 要求 CPU 至少空闲 20%
            memory_limit: 1024            # 要求内存至少剩余 1024 MB
            gpu_availability_limit: 30     # 要求 GPU 至少空闲 20%
            gpu_memory_limit: 2048         # 要求 GPU 显存至少空闲 2G
            min_available_threads_num: 1
        ********************************
        当前机器剩余值：
            cpu_availability = 40 （CPU 还有 40% 空闲）
            memory_available = 3072 （内存还有 3GB 可用）
            cpu_availability_limit = 20
            memory_limit = 1024
        
        CPU 维度: 每20%空闲支持1个并发：
            int((cpu_availability - self.cpu_availability_limit) / self.cpu_availability_limit
            cpu_concurrent = (40 - 20) / 20 = 1 → int(1) = 1

        内存维度：
            int(((memory_available - self.memory_limit)) / self.memory_limit)
            mem_concurrent = (3072 - 1024) / 1024 ≈ 2 → int(2) = 2

        # 取更严格的那个（瓶颈）
            min(1, 2) = 1
        ********************************
        cpu_availability = 15（低于阈值20）
        memory_available = 500（低于1024）
        min(-1, -1) = -1
        ********************************
        
        """
        cpu_result = min(int((cpu_availability - self.cpu_availability_limit) / self.cpu_availability_limit), int(((memory_available - self.memory_limit)) / self.memory_limit))
        
        # 负数情况：资源已经超标，加上负数，相当于减线程
        if cpu_result < 0:
            cpu_result = used_threads_num + cpu_result
        else:
            # 资源还够用  
            # 新建议 = 已用 + 可新增  used_threads_num = 2  cpu_result = 1（还能加1个）   建议总数 = 2 + 1 = 3
            cpu_result += used_threads_num

        if self.device == 'cpu':
            # cpu
            if cpu_result > self.min_available_threads_num:
                return min(need_threads_num, cpu_result)
            else:
                return self.min_available_threads_num
        else:
            # gpu  需要同时判断cpu
            gpu_availability, gpu_memory_available = DeviceChecker.check_gpu_availability()
            # 计算空余GPU资源
            gpu_result = min(int((gpu_availability - self.gpu_availability_limit) / self.gpu_availability_limit),
                             int((gpu_memory_available - self.gpu_memory_limit) / self.gpu_memory_limit),
                             cpu_result)
            if gpu_result < 0:
                gpu_result = used_threads_num + gpu_result
            else:
                gpu_result += used_threads_num
            if gpu_result > self.min_available_threads_num:
                return min(need_threads_num, gpu_result)
            else:
                return self.min_available_threads_num

@Singleton
class Pool:
    """
    ************线程池核心逻辑举例：*************
    1 用户提交了 5 个任务 max_threads = 5
    2 初始时 watch_dog 建议并发 3
    3 token_queue 有 3 个 1
    4 _task_schedule 拿到 3 个令牌，启动 3 个任务
    5 剩下 2 个任务在 task_queue 排队
    6 1 个任务完成 → token_queue 回收 1 个令牌
    7 _task_schedule 拿到这个令牌 → 启动第 4 个任务
    8 又完成一个 → 回收令牌 → 启动第 5 个任务
    ✅ 自动实现了“最多并发 3 个，做完一个再启动一个”


    ****************Pool 的核心组件************
    task_queue	任务排队的地方
    token_queue	控制并发的“许可证池”
    using_queue	记录哪些任务正在跑
    _task_schedule	调度员：拿证 → 取任务 → 执行
    adjust_threads	管理员：每 5 秒问一次“能并发几个？”
    JobWatchDog	决策者：根据资源决定并发上限
    _run_task	执行者：运行函数，完成后归还许可证`
    """
    def __init__(self, max_threads,watch_dog:JobWatchDog):
        
        if hasattr(self, 'initialized'):
            return
        self.initialized = True
        self.thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=max_threads)# 标准线程池
        self.task_queue = queue.Queue() # 任务排队的地方
        self.max_threads = max_threads
        # 可用令牌
        self.token_queue = queue.Queue(maxsize=max_threads)  # 控制并发的“许可证池”
        # 正在使用令牌
        self.using_queue = queue.Queue() # 记录哪些任务正在跑
        self.adjust_lock = threading.Lock()
        self.watch_dog = watch_dog
        self._initial_threads()

    def _initial_threads(self):
        """
        初始化令牌，要启动线程池，先问一下 JobWatchDog：现在系统资源够不够？我最多能给多少个‘工作许可证’？”
        """
        for _ in range(self.watch_dog.get_available_threads_num(self.max_threads, 0, 0)):
            self.token_queue.put(1)  # 初始化令牌，假设get_available_threads_num返回3个可用，则往token_queue里放3个令牌
        # 启动两个守护线程daemon
        # 线程池调整任务
        threading.Thread(target=self.adjust_threads, daemon=True).start() # 【管理员】 定期检查系统资源，决定要不要增加或减少“工作许可证”数量
        # 线程调度
        threading.Thread(target=self._task_schedule, daemon=True).start() # 【调度员】 负责从任务队列取任务，发许可证，让任务执行（拿证 → 取任务 → 执行）


    def submit_task(self, func, *args, **kwargs):
        """
        提交任务到任务队列，注意，此时只是请求，我要执行一个函数，只是进了队列
        """
        future = concurrent.futures.Future()
        task = (future, func, args, kwargs)
        self.task_queue.put(task)
        return future


    def _task_schedule(self):
        """
        守护线程：调度员  
        相当于一个“调度员”一直站着等任务。
        从令牌池，拿证 → 取任务 → 执行
        """
        while True:
            # 执行任务
            token = self.token_queue.get()  # 1. 拿一个许可证,会阻塞，直到有许可证可用。
            self.using_queue.put(token) # 2. 记录：这个证正在用
            future, func, args, kwargs = self.task_queue.get() # 3. 取一个任务,也是阻塞的，如果没有任务，就等着。
            self.thread_pool.submit(self._run_task, future, func, args, kwargs) # 4. 执行

    def adjust_threads(self):
        """
        守护线程：管理员
        管理员只管理令牌池，定期去问下，然后令牌池得扩容，缩容
        """
        while True:
            try:
                with self.adjust_lock:
                    # 动态调整
                    using_num = self.using_queue.qsize() # 正在运行的任务数
                    now_num = using_num + self.token_queue.qsize()  # 当前总并发数【正在跑的 + 可以跑的（空闲令牌）】
                    adjust_num = self.watch_dog.get_available_threads_num(self.max_threads, using_num, now_num) #JobWatchDog 会检查 CPU、内存、GPU 是否足够,返回一个建议值
                    print(f"调整线程数{adjust_num}")
                    if adjust_num > now_num: # 扩容: 建议 6，当前 4，多放2个	
                        for _ in range(adjust_num - now_num):
                            self.token_queue.put(1)
                    elif adjust_num < now_num: #缩容： 建议 2，当前 4	 从 token_queue 拿走 2 个令牌 → 新任务排队
                        # 缓慢释放
                        # 假设：
                        # 当前系统负载高 → 线程数多。
                        # 负载下降 → 监控判断可以缩容 → 一次性删掉 5 个令牌。
                        # 但下一个周期又有新任务进来 → 资源又够 → 又要扩容。
                        # 扩容 → 缩容 → 扩容 → 缩容……形成震荡。
                        # 这种频繁调整会：

                        # 增加调度开销。
                        # 导致任务执行不稳定。
                        # 可能误判资源状态（瞬时波动）。
                        self.token_queue.get()
                        continue
                time.sleep(5)
            except:
                pass

    def _run_task(self, future, func, args, kwargs):
        """"
        真正执行
        """
        try:
            result = func(*args, **kwargs)
            future.set_result(result)
        except Exception as e:
            future.set_exception(e)
        finally:
            # 释放令牌
            self.token_queue.put(1)
            self.using_queue.get()
